# /usr/bin/python3
#-*-coding:utf-8-*-
# 运行环境： ubuntu 16.04 LTS
# tensorflow 1.13.1
# opencv 3.4.1.15
# 数据集： COCO 2017
# 查看COCO数据集中的类的样本数, 以及是否做了标定框
import cv2
import matplotlib as mt
import matplotlib.pyplot as plt
import tensorflow as tf
import numpy as np
import os
import time
import PIL
import skimage.io as io
import xml.etree.ElementTree as ET

from pycocotools.coco import COCO
# 这里顺序要放在最后, 不然这个界面显示的GUI不会生效
mt.use('TkAgg');
# 原数据集标注信息文件路径
annfile = r'COCO/dataset_training/annotations_trainval2017/annotations/instances_train2017.json';
coco = COCO(annfile);
# 本执行文件与COCO目录是在同一目录下, path路径下全是图片
path = 'COCO//dataset_training//train2017//train2017//';
# 标注信息有多个标定框(同一张图片的同一个类别有多个)
picture_count = 0;
# 记录自己一共写入了多少条样本数据
winder_sum = 0;
# 标定框的坐标值保留小数的位数
Decimal_digits = 3;
# car           的ID号是 3;    数量： 12251
# backpack      的ID号是 27;   数量： 5528
# tie           的ID号是 32;   数量： 3810
# handbag       的ID号是 31;   数量： 6841
# suitcase      的ID号是 33;   数量： 2402
# bottle        的ID号是 44;   数量： 8501
# wine glass    的ID号是 46;   数量： 2533
# cup           的ID号是 47;   数量： 9189
# knife         的ID号是 49;   数量： 4326
# spoon         的ID号是 50;   数量： 3529
# bowl          的ID号是 51;   数量： 7111
# chair         的ID号是 62;   数量： 12774
# couch         的ID号是 63;   数量： 4423
# potted plant  的ID号是 64;   数量： 4452
# bed           的ID号是 65;   数量： 3682
# dining table  的ID号是 67;   数量： 11837
# toilet        的ID号是 70;   数量： 3353
# banana        的ID号是 52;   数量： 2243
# apple         的ID号是 53;   数量： 1586
# sandwich      的ID号是 54;   数量： 2365
# orange        的ID号是 55;   数量： 1699
# hot dog       的ID号是 58;   数量： 1222
# pizza         的ID号是 59;   数量： 3166
# cake          的ID号是 61;   数量： 2925
# tv            的ID号是 72;   数量： 4561
# laptop        的ID号是 73;   数量： 3524
# mouse         的ID号是 74;   数量： 1876
# keyboard      的ID号是 76;   数量： 2115
# cell phone    的ID号是 77;   数量： 4803
# book          的ID号是 84;   数量： 5332
# clock         的ID号是 85;   数量： 4659
# vase          的ID号是 86;   数量： 3593
# scissors      的ID号是 87;   数量： 947
# microwave     的ID号是 78;   数量： 1547
# oven          的ID号是 79;   数量： 2877
# refrigerator  的ID号是 82;   数量： 2360
# bicycle       的ID号是 2;    数量： 3252
# motorcycle    的ID号是 4;    数量： 3502
# train         的ID号是 7;    数量： 3588
# truck         的ID号是 8;    数量： 6127
# boat          的ID号是 9:    数量： 3025
# cat           的ID号是 17;   数量： 4114
# dog           的ID号是 18;   数量： 4385
# bird          的ID号是 16;   数量： 3237
# traffic light 的ID号是 10;   数量： 4139
# bench         的ID号是 15;   数量： 5570
# umbrella      的ID号是 28;   数量： 3968
# person        的ID号是 1;    数量： 64115
# bus           的ID号是 6;    数量： 3952

# 创建一个类的名称的列表
class_name = ['car', 'backpack', 'tie', 'handbag', 'suitcase', 'bottle', 'wine glass', 'cup', 'knife', 'spoon',
              'bowl', 'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'banana', 'apple', 'sandwich',
              'orange', 'hot dog', 'pizza', 'cake', 'tv', 'laptop', 'mouse', 'keyboard', 'cell phone', 'book', 'clock',
              'vase', 'scissors', 'microwave', 'oven', 'refrigerator', 'bicycle', 'motorcycle', 'train', 'truck', 'boat',
              'cat', 'dog', 'bird', 'traffic light', 'bench', 'umbrella', 'bus'];
print(len(class_name));
temp_class_name = class_name;
# 用于剔除COCO数据集中不存在的标签
for i, name in enumerate(class_name):
    catIds = coco.getCatIds(catNms=name);
    if catIds == [] :
        temp_class_name.remove(name);
        print('这个 %s 不存在！' % name);
print(temp_class_name);
class_name = temp_class_name;
print(len(class_name));


# 开始创建文件, 用于把所有的图片路径以及标定框写入TXT 文本中;
# 用于存放整理出来的数据集的文本的文件夹名字;
datasets_folder = 'winder_datasets';
if not os.path.exists(datasets_folder):
    os.makedirs(datasets_folder);
train_dataset = os.path.join(datasets_folder, 'train of winder');
print(train_dataset); # winder_datasets/train of winder
# 对于一个即将要写入数据集信息的文本, 命名一个新的名字
datasets_file_name = 'winder_COCO';
label_file_name = 'winder_label';
type_datasets = 'txt';


# 训练集label文本的路径
train_label_path = train_dataset + '/' + label_file_name + '.' + type_datasets;
if not os.path.exists(train_label_path):
    #os.makedirs(train_dataset_path);
    f = open(train_label_path, 'w');
    f.close();

names = {};
with open(train_label_path, 'a', encoding='utf-8') as f:
    for i in range(len(class_name)):
        temp_name = class_name[i];
        names[temp_name] = i;
        label_input = str(i) + '  ' + temp_name;
        f.write(label_input + '\n');
    print(names);

# 训练集文本的路径
train_dataset_path = train_dataset + '/' + datasets_file_name + '.' + type_datasets;
print('训练数据集的文本的路径：', train_dataset_path);
if not os.path.exists(train_dataset_path):
    #os.makedirs(train_dataset_path);
    f = open(train_dataset_path, 'w');
    f.close();

with open(train_dataset_path, 'a', encoding='utf-8') as f:
    for names_1 in class_name:
        # 得到数据集中类别的ID号
        catIds = coco.getCatIds(catNms=names_1);
        # 得到所有这个类的图片的文件名
        imgIds = coco.getImgIds(catIds=catIds[0]);
        # 所属类别的图片的数量
        print('这个 %s 名字的类别的样本数量：' % names_1, end=' ');
        print(len(imgIds), end='\n');
        winder_sum += len(imgIds);
        for ID in range(len(imgIds)):
            # img 这个变量的结构是外面是列表, 列表中嵌套字典
            img = coco.loadImgs(imgIds[ID]);
            '''
            img = coco.loadImgs(imgIds[2500]);
            print(img);
            
            [{'license': 2, 'file_name': '000000203235.jpg', 
            'coco_url': 'http://images.cocodataset.org/train2017/000000203235.jpg', 
            'height': 640, 'width': 427, 'date_captured': '2013-11-18 17:28:36', 
            'flickr_url': 'http://farm3.staticflickr.com/2135/2059040557_a69e5f4121_z.jpg', 
            'id': 203235}]
            '''
            # print(img[0]['file_name']) 输出： 000000203235.jpg
            img_path = path + str(img[0]['file_name']);
            if not os.path.exists(img_path):
                winder_sum -= 1;
                continue;
            # 开始读取标注信息
            annIds = coco.getAnnIds(imgIds=img[0]['id'], catIds=catIds, iscrowd=None);
            anns = coco.loadAnns(annIds);
            annotation = anns[0]['bbox'];
            # 记录一下一张图的一个类有多个标定框
            if len(annotation) > 4:
                picture_count += 1;
                print('类别的名字： ', names_1);
                print('图片的名字： ', str(img[0]['file_name']));
            xmin, ymin, xmax, ymax = annotation[0], annotation[1], \
                                     annotation[0] + annotation[2], annotation[1] + annotation[3];
            xmin, ymin, xmax, ymax = round(xmin, Decimal_digits), round(ymin, Decimal_digits), \
                                     round(xmax, Decimal_digits), round(ymax, Decimal_digits);
            Class_ID = names[names_1];
            '''
            注意啊, 只要是iscrowd=0那么segmentation就是polygon格式；只要iscrowd=1那么segmentation就是RLE格式.
            另外, 每个对象（不管是iscrowd=0还是iscrowd=1）都会有一个矩形框bbox, 矩形框左上角的坐标和矩形框的长宽会以数组的形式提供, 
            数组第一个元素就是左上角的横坐标值。
            '''
            text_input = img_path + ' ' + ','.join([str(xmin), str(ymin),
                                                str(xmax), str(ymax),
                                                str(Class_ID)]);
            f.write(text_input + '\n');
    print('有多少个样本中同一个类的标定框有多个： ', picture_count);
    print('所有存在并且写入文本的样本数据总数： ', winder_sum);

# 输出LOG ：
'''
{'traffic light': 44, 'mouse': 26, 'cell phone': 28, 'tie': 2, 'banana': 17, 
'cup': 7, 'bus': 47, 'bench': 45, 'keyboard': 27, 'bowl': 10, 'oven': 34, 'handbag': 3, 
'motorcycle': 37, 'orange': 20, 'umbrella': 46, 'bottle': 5, 'hot dog': 21, 'microwave': 33, 
'toilet': 16, 'scissors': 32, 'tv': 24, 'car': 0, 'boat': 40, 'pizza': 22, 'refrigerator': 35, 
'backpack': 1, 'sandwich': 19, 'potted plant': 13, 'knife': 8, 'wine glass': 6, 'dining table': 15, 
'spoon': 9, 'clock': 30, 'laptop': 25, 'dog': 42, 'apple': 18, 'chair': 11, 'bird': 43, 'couch': 12, 
'cat': 41, 'book': 29, 'suitcase': 4, 'vase': 31, 'truck': 39, 'bicycle': 36, 'cake': 23, 'bed': 14, 'train': 38}

训练数据集的文本的路径： winder_datasets/train of winder/winder_COCO.txt

这个 car 名字的类别的样本数量： 12251
这个 backpack 名字的类别的样本数量： 5528
这个 tie 名字的类别的样本数量： 3810
这个 handbag 名字的类别的样本数量： 6841
这个 suitcase 名字的类别的样本数量： 2402
这个 bottle 名字的类别的样本数量： 8501
这个 wine glass 名字的类别的样本数量： 2533
这个 cup 名字的类别的样本数量： 9189
这个 knife 名字的类别的样本数量： 4326
这个 spoon 名字的类别的样本数量： 3529
这个 bowl 名字的类别的样本数量： 7111
这个 chair 名字的类别的样本数量： 12774
这个 couch 名字的类别的样本数量： 4423
这个 potted plant 名字的类别的样本数量： 4452
这个 bed 名字的类别的样本数量： 3682
这个 dining table 名字的类别的样本数量： 11837
这个 toilet 名字的类别的样本数量： 3353
这个 banana 名字的类别的样本数量： 2243
这个 apple 名字的类别的样本数量： 1586
这个 sandwich 名字的类别的样本数量： 2365
这个 orange 名字的类别的样本数量： 1699
这个 hot dog 名字的类别的样本数量： 4385
这个 pizza 名字的类别的样本数量： 3166
这个 cake 名字的类别的样本数量： 2925
这个 tv 名字的类别的样本数量： 4561
这个 laptop 名字的类别的样本数量： 3524
这个 mouse 名字的类别的样本数量： 1876
这个 keyboard 名字的类别的样本数量： 2115
这个 cell phone 名字的类别的样本数量： 4803
这个 book 名字的类别的样本数量： 5332
这个 clock 名字的类别的样本数量： 4659
这个 vase 名字的类别的样本数量： 3593
这个 scissors 名字的类别的样本数量： 947
这个 microwave 名字的类别的样本数量： 1547
这个 oven 名字的类别的样本数量： 2877
这个 refrigerator 名字的类别的样本数量： 2360
这个 bicycle 名字的类别的样本数量： 3252
这个 motorcycle 名字的类别的样本数量： 3502
这个 train 名字的类别的样本数量： 3588
这个 truck 名字的类别的样本数量： 6127
这个 boat 名字的类别的样本数量： 3025
这个 cat 名字的类别的样本数量： 4114
这个 dog 名字的类别的样本数量： 4385
这个 bird 名字的类别的样本数量： 3237
这个 traffic light 名字的类别的样本数量： 4139
这个 bench 名字的类别的样本数量： 5570
这个 umbrella 名字的类别的样本数量： 3968
这个 bus 名字的类别的样本数量： 3952

有多少个样本中同一个类的标定框有多个：  0
所有存在并且写入文本的样本数据总数：  211964
'''